/**
 * 并查集（Union-Find）算法
 */
class UnionFind {
  // n 为图中节点的个数
  constructor(n) {
    // 连通分量个数
    this._count = n
    // 存储每个节点的父节点
    this.parent = Array.from({ length: n }, (_, index) => index)
    // 存储每个节点的权重（以该节点为根的树中节点个数）
    this.size = new Array(n).fill(1)
  }

  // 将节点 p 和节点 q 连通
  connect(p, q) {
    const rootP = this.find(p)
    const rootQ = this.find(q)

    if (rootP === rootQ) {
      return false
    }

    // 小树接到大树下面，较平衡
    if (this.size[rootP] < this.size[rootQ]) {
      this.parent[rootP] = rootQ
      this.size[rootQ] += this.size[rootP]
    } else {
      this.parent[rootQ] = rootP
      this.size[rootP] += this.size[rootQ]
    }
    // 两个连通分量合并成一个连通分量
    this._count--
    return true
  }

  // 判断节点 p 和节点 q 是否连通
  isConnected(p, q) {
    const rootP = this.find(p)
    const rootQ = this.find(q)
    return rootP === rootQ
  }

  // 返回图中的连通分量个数
  count() {
    return this._count
  }

  // 返回 x 所在的连通分量的大小
  componentSizeOf(x) {
    const rootX = this.find(x)
    return this.size[rootX]
  }
  
  // 返回节点 x 的根节点
  find(x) {
    if (this.parent[x] !== x) {
      this.parent[x] = this.find(this.parent[x])
    }
    return this.parent[x]
  }
}

// test
const uf = new UF(5)
uf.connect(1, 2)
uf.connect(3, 2)
uf.connect(4, 2)
console.log(uf.isConnected(1, 4)) // true
console.log(uf.count()) // 2
console.log(uf.isConnected(1, 0)) // false
console.log(uf.isConnected(1, 3)) // true
console.log(uf.isConnected(0, 4)) // false
console.log(uf.isConnected(0, 3)) // false
console.log(uf.componentSizeOf(1)) // 4
console.log(uf.componentSizeOf(0)) // 1
console.log(uf.componentSizeOf(3)) // 4